/*
 * Copyright (c) 2021-2024 Huawei Device Co., Ltd.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "bpf_loader.h"
#include "i_netfirewall_callback.h"
#include "netfirewall_callback_stub.h"
#include <arpa/inet.h>
#include <gtest/gtest.h>
#include <netinet/in.h>
#include <sys/socket.h>
#include <unistd.h>

#define private public
#define protected public
#include "bitmap_manager.h"
#include "bpf_netfirewall.h"

#define FILE_NAME (strrchr(__FILE__, '/') + 1)
#define MANUAL_TEST 0

using namespace std;
using namespace testing::ext;
using namespace OHOS;
using namespace OHOS::NetManagerStandard;

static sptr<NetFirewallIpRule> GeIpFirewallRule(NetFirewallRuleDirection dir, string addr)
{
    sptr<NetFirewallIpRule> rule = (std::make_unique<NetFirewallIpRule>()).release();
    if (!rule) {
        return rule;
    }
    rule->ruleAction = FirewallRuleAction::RULE_DENY;
    rule->appUid = 0;

    std::vector<NetFirewallIpParam> remoteIp;
    NetFirewallIpParam remoteIpParam;
    remoteIpParam.family = 1;
    remoteIpParam.type = 1;
    remoteIpParam.address = addr;
    remoteIpParam.mask = 0;
    remoteIpParam.startIp = "";
    remoteIpParam.endIp = "";
    remoteIp.push_back(remoteIpParam);
    rule->remoteIps = remoteIp;

    rule->protocol = NetworkProtocol::ICMP;

    return rule;
}

class NetsysNetFirewallTest : public testing::Test {
public:
    static void SetUpTestCase();
    static void TearDownTestCase();
    void SetUp();
    void TearDown();
};

void NetsysNetFirewallTest::SetUpTestCase()
{
    shared_ptr<NetsysBpfNetFirewall> bpfNetFirewall = NetsysBpfNetFirewall::GetInstance();
    if (!bpfNetFirewall->IsBpfLoaded()) {
        auto ret = OHOS::NetManagerStandard::LoadElf(FIREWALL_BPF_PATH);
        printf("LoadElf is %d\n", ret);

        if (ret == ElfLoadError::ELF_LOAD_ERR_NONE) {
            bpfNetFirewall->SetBpfLoaded(true);
        }
        bpfNetFirewall->StartListener();
    }
}

void NetsysNetFirewallTest::TearDownTestCase()
{
    shared_ptr<NetsysBpfNetFirewall> bpfNetFirewall = NetsysBpfNetFirewall::GetInstance();
    if (bpfNetFirewall->IsBpfLoaded()) {
        bpfNetFirewall->StopListener();
        auto ret = OHOS::NetManagerStandard::UnloadElf(FIREWALL_BPF_PATH);
        printf("UnloadElf is %d\n", ret);
        if (ret == ElfLoadError::ELF_LOAD_ERR_NONE) {
            bpfNetFirewall->SetBpfLoaded(false);
        }
    }
}

void NetsysNetFirewallTest::SetUp() {}

void NetsysNetFirewallTest::TearDown() {}

static void PrintBitmap(RuleCode &code, const char *reason)
{
    printf("%s: ", reason);
    for (int i = 0; i < BITMAP_LEN; i++) {
        printf("%X", code.val[i]);
    }
    printf("\n");
}

static void Trace()
{
    printf("Trace: [%s] line: %d\n", FILE_NAME, __LINE__);
}

class TestNetFirewallCallbackStub : public OHOS::NetsysNative::NetFirewallCallbackStub {
public:
    int32_t OnIntercept(OHOS::sptr<InterceptRecord> &info)
    {
        if (!info) {
            return -1;
        }

        std::thread t = std::thread([&]() {
            printf("\ttransProtocol=%u\n", info->protocol);
            printf("\tsourcePort=%u\n", info->sourcePort);
            printf("\tdestPort=%u\n", info->destPort);
            printf("\tsourceIp=%s\n", (info->sourceIp).c_str());
            printf("\tdestIp=%s\n", (info->destIp).c_str());
            printf("\tappUid=%d\n", info->appUid);
        });
        t.join();

        return 0;
    }
};

static bool g_keepWait = true;

void SignalHandler(int signum)
{
    g_keepWait = false;
}

static void Wait()
{
#if MANUAL_TEST
    signal(SIGINT, SignalHandler);
    signal(SIGQUIT, SignalHandler);
    signal(SIGABRT, SignalHandler);
    signal(SIGKILL, SignalHandler);
    signal(SIGSTOP, SignalHandler);

    g_keepWait = true;
    const unsigned int TIME = 3;
    while (g_keepWait) {
        sleep(TIME);
    }
#endif
}

HWTEST_F(NetsysNetFirewallTest, NetsysNetFirewallTest001, TestSize.Level0)
{
    Trace();
    shared_ptr<NetsysBpfNetFirewall> bpfNetFirewall = NetsysBpfNetFirewall::GetInstance();

    OHOS::sptr<OHOS::NetsysNative::INetFirewallCallback> callback = new (nothrow)TestNetFirewallCallbackStub;
    int32_t ret = bpfNetFirewall->RegisterCallback(callback);
    EXPECT_EQ(ret, 0);

    ret = bpfNetFirewall->UnregisterCallback(callback);
    EXPECT_EQ(ret, 0);

    Ip4Key srcIp = 0;
    string src = "192.168.8.116";
    bpfNetFirewall->IpAtoN(AF_INET, src, &srcIp);

    Ip4Key dstIp = 0;
    string dst = "192.168.1.5";
    bpfNetFirewall->IpAtoN(AF_INET, dst, &dstIp);
    InterceptEvent interceptEv = {
        .dir = INGRESS,
        .family = AF_INET,
        .protocol = IPPROTO_TCP,
        .sport = 80,
        .dport = 5684,
        .appuid = 0,
    };
    interceptEv.ipv4.saddr = srcIp;
    interceptEv.ipv4.daddr = dstIp;

    bpfNetFirewall->NotifyInterceptEvent(&interceptEv);

    Event ev = {
        .intercept = interceptEv,
        .type = EVENT_INTERCEPT,
        .len = sizeof(InterceptEvent),
    };
    ret = NetsysBpfNetFirewall::HandleEvent(NULL, &ev, sizeof(ev));
    EXPECT_EQ(ret, 0);

    DebugEvent debugEv = {
        .type = DBG_GENERIC,
        .dir = INGRESS,
        .arg1 = 0,
    };
    NetsysBpfNetFirewall::HandleDebugEvent(&debugEv);
}

HWTEST_F(NetsysNetFirewallTest, NetsysNetFirewallTest002, TestSize.Level0)
{
    Trace();
    Bitmap a(1);
    Bitmap b(a);

    EXPECT_TRUE(a.IsValid());
    b.Clear();
    b.Set(3);
    EXPECT_TRUE(a.Hash() != 0);
    EXPECT_TRUE(a.SpecialHash() != 0);
    EXPECT_TRUE(b.Get(3) != 0);
    EXPECT_TRUE(a.CountBits() > 0);

    EXPECT_TRUE(b.Get() != nullptr);

    a.Or(b);
    a.And(b);
    b.Not(0);

    EXPECT_FALSE(a == b);

    Trace();
    EXPECT_TRUE(b.ToString() != nullptr);
}

HWTEST_F(NetsysNetFirewallTest, NetsysNetFirewallTest003, TestSize.Level0)
{
    Trace();
    BpfUnorderedMap<int> map;
    int key = 1;
    Bitmap val(10);
    Bitmap other(20);
    map.OrInsert(key, val);
    map.OrInsert(key + 1, val);
    map.OrForEach(other);

    map.Delete(key);
    EXPECT_FALSE(map.Empty());
    map.Clear();
    EXPECT_TRUE(map.Empty());

    EXPECT_TRUE(map.Get().empty());
}

HWTEST_F(NetsysNetFirewallTest, NetsysNetFirewallTest004, TestSize.Level0)
{
    Trace();
    IpV6Helper helper;
    IpV6Helper other;
    string ip = "2001:0db8:85a3:0000:0000:8a2e:0370:7334";

    helper.Init(ip);
    helper.Init(ip, 12);

    helper.Increase();
    EXPECT_FALSE(helper.IsLess(other));

    string out;
    helper.GetStr(out);
    EXPECT_TRUE(out.length() > 0);

    int32_t ret = helper.ParseStr(ip);
    EXPECT_EQ(ret, 0);
}

HWTEST_F(NetsysNetFirewallTest, NetsysNetFirewallTest005, TestSize.Level0)
{
    Trace();
    std::vector<NetFirewallIpRule> ruleList;
    sptr<NetFirewallIpRule> rule = GeIpFirewallRule(NetFirewallRuleDirection::RULE_IN, "153.3.238.110");
    ruleList.push_back(*rule);
    rule->remoteIps.front().address = "153.3.238.102";
    ruleList.push_back(*rule);

    BitmapManager manager;
    int ret = manager.BuildBitmapMap(ruleList);
    EXPECT_EQ(ret, 0);

    EXPECT_FALSE(manager.GetSrcIp4Map().empty());
    EXPECT_FALSE(manager.GetSrcIp6Map().empty());
    EXPECT_FALSE(manager.GetDstIp4Map().empty());
    EXPECT_FALSE(manager.GetDstIp6Map().empty());
    EXPECT_FALSE(manager.GetSrcPortMap().Empty());
    EXPECT_FALSE(manager.GetDstPortMap().Empty());
    EXPECT_FALSE(manager.GetProtoMap().Empty());
    EXPECT_FALSE(manager.GetActionMap().empty());
    EXPECT_FALSE(manager.GetAppIdMap().Empty());

    rule->ruleDirection = NetFirewallRuleDirection::RULE_OUT;
    ruleList.push_back(*rule);
    rule->remoteIps.front().address = "153.3.238.103";
    ruleList.push_back(*rule);

    manager.BuildBitmapMap(ruleList);
    EXPECT_EQ(ret, 0);

    manager.BuildBitmapMap(ruleList);
    manager.BuildMarkBitmap(ruleList);
    manager.BuildNoMarkBitmap(ruleList);
    Bitmap a(1);
    manager.InsertIpBitmap(rule->remoteIps, true, a);
    a.Set(5);
    manager.InsertIpBitmap(rule->remoteIps, false, a);

    vector<NetFirewallPortParam> ports;
    NetFirewallPortParam sport;
    sport.startPort = 80;
    sport.endPort = 80;
    ports.emplace_back(sport);

    manager.OrInsertPortBitmap(ports, a, manager.GetSrcPortMap());

    bool res = manager.IsNotNeedPort(rule->protocol);
    EXPECT_TRUE(res);
}

HWTEST_F(NetsysNetFirewallTest, NetsysNetFirewallTest006, TestSize.Level0)
{
    Trace();
    std::vector<NetFirewallIpRule> ruleList;
    sptr<NetFirewallIpRule> rule = GeIpFirewallRule(NetFirewallRuleDirection::RULE_IN, "153.3.238.110");
    ruleList.push_back(*rule);
    rule->remoteIps.front().address = "153.3.238.102";
    ruleList.push_back(*rule);

    Trace();
    shared_ptr<NetsysBpfNetFirewall> bpfNetFirewall = NetsysBpfNetFirewall::GetInstance();
    bool ret = bpfNetFirewall->SetFirewallRule(ruleList, NetFirewallRuleDirection::RULE_IN);
    EXPECT_TRUE(ret);

    vector<pair<Ip4Key, RuleCode>> vec;
    int res = bpfNetFirewall->ReadBpfMap(GET_MAP_PATH(true, saddr), vec);
    EXPECT_TRUE(res == 0);
    EXPECT_TRUE(vec.size() > 0);
    printf("\n%s:\n", GET_MAP_PATH(true, saddr));
    for (auto pair : vec) {
        Ip4Key k = pair.first;
        RuleCode v = pair.second;

        struct in_addr in;
        in.s_addr = k;
        printf("key: %s ", inet_ntoa(in));
        PrintBitmap(v, "value");
    }
    Wait();
}